/* Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This file is a part of the CANN Open Software.
 * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 * ===================================================================================================================*/
#ifndef METADEF_BASE_CONTEXT_BUILDER_OP_INFO_IMPL_H_
#define METADEF_BASE_CONTEXT_BUILDER_OP_INFO_IMPL_H_
#include <memory>
#include <vector>
#include "types.h"
#include "graph/ascend_string.h"
#include "exe_graph/runtime/storage_shape.h"
#include "exe_graph/runtime/range.h"
#include "exe_graph/runtime/shape.h"
#include "graph/any_value.h"
#include "common/checker.h"
#include "exe_graph/runtime/tensor.h"
#include "exe_graph/runtime/kernel_context.h"
#include "base/context_builder/op_info.h"
using namespace ge;

namespace gert {
struct ContextTensorDesc {
  ge::DataType dtype_ = ge::DataType::DT_MAX;
  ge::Format origin_format_ = ge::Format::FORMAT_MAX;
  ge::Format storage_format_ = ge::Format::FORMAT_MAX;
  gert::StorageShape shape_ = {};
};
struct OpInfoImpl {
  OpInfoImpl() = default;
  ~OpInfoImpl() = default;

  void SetOpType(const AscendString &op_type) {
    op_type_ = op_type.GetString();
  }
  void SetOpName(const AscendString &op_name) {
    op_name_ = op_name.GetString();
  }
  void SetIONum(size_t input_num, size_t output_num) {
    if (!input_instance_.empty() || !output_instance_.empty()) {
      GELOGW("IO has been set. Set IO Num failed!");
      return;  // 已经设置过输入输出， 无需不允许再次设置
    }
    input_num_ = input_num;
    output_num_ = output_num;
    input_tensor_descs_.resize(input_num_);
    output_tensor_descs_.resize(output_num_);
    input_instance_.resize(input_num_, 1);
    output_instance_.resize(output_num_, 1);
    input_instance_num_ = input_num_;
    output_instance_num_ = output_num_;
  }
  void SetIOInstanceNum(const std::vector<uint32_t> &input_instance, const std::vector<uint32_t> &output_instance) {
    input_instance_ = input_instance;
    output_instance_ = output_instance;
    input_num_ = input_instance.size();
    output_num_ = output_instance.size();
    input_instance_num_ = 0U;
    output_instance_num_ = 0U;
    for (const auto &num : input_instance) {
      input_instance_num_ += num;
    }
    input_tensor_descs_.resize(input_instance_num_);
    for (const auto &num : output_instance) {
      output_instance_num_ += num;
    }
    output_tensor_descs_.resize(output_instance_num_);
  }
  void SetInputTd(size_t index, ge::DataType dtype, ge::Format origin_format, ge::Format storage_format,
                  const gert::StorageShape &shape) {
    if (index >= input_tensor_descs_.size()) {
      GELOGE(ge::FAILED, "Input tensor index %zu exceeds size %zu", index, input_tensor_descs_.size());
      return;
    }
    auto td = ContextTensorDesc();
    td.dtype_ = dtype;
    td.origin_format_ = origin_format;
    td.storage_format_ = storage_format;
    td.shape_ = shape;
    input_tensor_descs_[index] = std::move(td);
  }
  void SetOutputTd(size_t index, ge::DataType dtype, ge::Format origin_format, ge::Format storage_format,
                   const gert::StorageShape &shape) {
    if (index >= output_tensor_descs_.size()) {
      GELOGE(ge::FAILED, "Output tensor index %zu exceeds size %zu", index, output_tensor_descs_.size());
      return;
    }
    auto td = ContextTensorDesc();
    td.dtype_ = dtype;
    td.origin_format_ = origin_format;
    td.storage_format_ = storage_format;
    td.shape_ = shape;
    output_tensor_descs_[index] = std::move(td);
  }
  template<typename AttrTypeT>
  void SetAttr(const AscendString &attr_name, AttrTypeT attr) {
    attrs_.emplace_back(
        std::pair<std::string, ge::AnyValue>(attr_name.GetString(), ge::AnyValue::CreateFrom<AttrTypeT>(attr)));
  }

  bool CheckParams() const {
    GE_ASSERT_TRUE(op_type_.size() != 0 && op_name_.size() != 0 && input_num_ != 0 && output_num_ != 0,
                   "Invalid params, op_type: %s, op_name: %s, input_num: %u, output_num: %u", op_type_.c_str(),
                   op_name_.c_str(), input_num_, output_num_);
    for (const auto &i : input_instance_) {
      GE_ASSERT_TRUE(i > 0, "Input instance num should be greater than 0, but got %u", i);
    }
    return true;
  }

  std::string op_type_;
  std::string op_name_;
  std::vector<std::pair<std::string, ge::AnyValue>> attrs_;  // 传递给context中的compute_node_info
  uint32_t input_num_ = 0U;
  uint32_t input_instance_num_ = 0U;
  uint32_t output_num_ = 0U;
  uint32_t output_instance_num_ = 0U;
  std::vector<uint32_t> input_instance_;
  std::vector<uint32_t> output_instance_;
  std::vector<ContextTensorDesc> input_tensor_descs_;   // 传递给context中的compute_node_info
  std::vector<ContextTensorDesc> output_tensor_descs_;  // 传递给context中的compute_node_info
};
class OpInfoHelper {
 public:
  static OpInfoImpl *GetPtr(const OpInfo &op_info) {
    return op_info.impl_.get();
  }
};
}  // namespace gert

#endif  // METADEF_BASE_CONTEXT_BUILDER_OP_INFO_IMPL_H_